import numpy as np
from flask import Flask
from flask import request
from flask_cors import CORS
from PIL import Image
import torchvision
import cv2 as cv
from torch.nn import Linear
import torch
import torchvision.transforms as transforms
istest = False
herbs_name = {}
labelspath = './resource/labels.txt'
with open(labelspath, 'r', encoding='utf-8') as f:
    for line in f:
        line = line.strip()
        parts = line.split(maxsplit=1)
        if len(parts) == 2:
            id, name = parts
            herbs_name[id] = name

m_mean = [0.6119, 0.5443, 0.4698]
m_std = [0.2263, 0.2417, 0.2556]
data_transform = transforms.Compose([
    transforms.Resize((299, 299)),
    transforms.CenterCrop(180),
    transforms.ToTensor(),
    transforms.Normalize(mean=m_mean, std=m_std)
])
device = torch.device("cuda")
modelpath = './models/chinese_mechicine_model_1.pth'

model = torchvision.models.vgg16(weights='DEFAULT')
model.classifier.add_module('7', Linear(1000, 163))
path = torch.load(modelpath, weights_only=True)
model.load_state_dict(path)
model.to(device)
if istest:
    print(model)

app = Flask(__name__)
CORS(app)
@app.route('/tools/herb', methods=['POST'])
def herb_classifiaction():
    # 获得发送过来的图片
    buf = request.files['image'].read()
    img = cv.imdecode(np.frombuffer(buf, dtype=np.uint8), flags=cv.IMREAD_COLOR)
    # 转换为PIL格式
    image = Image.fromarray(cv.cvtColor(img, cv.COLOR_BGR2RGB))
    # 处理数据
    timage = data_transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(timage)
    #result = output.argmax(1)
    result = {'code': 1, 'msg': None, 'data': {}}
    result['data']['name'] = herbs_name[str(output.argmax(1).tolist()[0])]
    return result

if __name__ == '__main__':
    app.run(debug=istest)

